from langchain.schema.runnable import RunnableMap,RunnableParallel,RunnableBranch,RunnableLambda
from langchain_vector_data import get_retriever
from langchain_study_deepseek import get_llm
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

def get_prompt_from_langchain():
    no_do_prompt_qu = """
    你现在是一名专业的中医药医生，通过历史问诊对话以及原始症状给出合适的药方。如果你不知道答案，就说你不知道，不要编造答案。请根据以下提供的问诊记录回答问题,如果没有提供历史问诊对话直接回答没有问诊记录，无法开药方。

    
    历史问诊记录: {diagnosis_chat_history}
    病人症状：{symptom}
    请按照以下格式返回回答:
    答案: <你的回答>
    药方: <答案必须包含相应的知识来源，用逗号分隔>
    开药时间：<当前时间年月日>
    """
    no_do_prompt = PromptTemplate(input_variables=["diagnosis_chat_history","symptom"],
                                 template=no_do_prompt_qu)
    prompt_template_question = """你现在是一名专业的中医药医生，通过历史问诊对话和病人症状给出合适的药方。如果你不知道答案，就说你不知道，不要编造答案。请根据以下提供的问诊记录回答问题,如果没有提供历史问诊对话直接回答没有问诊记录，无法开药方。

        历史问诊记录：
        {diagnosis_chat_history}
        病人症状：{symptom}
        请按照以下格式返回药方:
        药方: <你的回答>
        来源: <答案必须包含相应的知识来源，用逗号分隔>
        开药时间：<当前时间年月日>
        """
        # 定义两个必需的 prompt
    question_prompt = PromptTemplate(
        template=prompt_template_question,
        input_variables=["diagnosis_chat_history","symptom"]
    )
    prompt_template_refine = """你现在是一名专业的中医药医生，通过历史问诊对话给出合适的药方。现在你将获得更多的知识，请结合新的内容进一步优化回答。如果你不知道答案，就说你不知道，不要编造答案。
        病人症状：{symptom}
        新增上下文:
        {context_str}
        当前回答：
        {existing_answer}
        
        请按照以下格式返回回答:
        药方: <你的回答>
        来源: <答案必须包含相应的知识来源，用逗号分隔>
        开药时间：<当前时间年月>
        """
        
    refine_prompt = PromptTemplate(
        template=prompt_template_refine,
        input_variables=["existing_answer", "context_str","symptom"]
)

    return question_prompt,refine_prompt,no_do_prompt

def refine_chain_with_docs(symptom_dict,docs,initial_prompt,refine_prompt,fallback_prompt,llm):
    symptom = symptom_dict['symptom']
    if not docs:
        # 没有找到相关文档，使用备用 prompt
        answer = (fallback_prompt | llm | StrOutputParser()).invoke({"symptom":symptom,"diagnosis_chat_history":symptom_dict['diagnosis_chat_history']})
        return {"result": answer, "sources": []}

    # 初步回答
    context = docs[0].page_content
    answer = (initial_prompt | llm | StrOutputParser()).invoke({
        "context_str": context,
        "symptom": symptom
    })

    # 逐个 refine
    for doc in docs[1:]:
        answer = (refine_prompt | llm | StrOutputParser()).invoke({
            "existing_answer": answer,
            "context_str": doc.page_content,
            "symptom": symptom
        })

    return {"result": answer, "sources": docs}

def get_LECL_chain():
    retriever = get_retriever()
    llm = get_llm()
    question_prompt,refine_prompt,no_do_prompt = get_prompt_from_langchain()
    # prepare_input = RunnableLambda(lambda q: {"symptom": q})
    # 调用转换器
    retrieval_chain = RunnableLambda(lambda x: {**x, "docs": retriever.invoke(x["symptom"])})

    final_chain = (retrieval_chain | RunnableLambda(lambda x:refine_chain_with_docs(x,x["docs"],question_prompt,refine_prompt,no_do_prompt,llm))
    )
    return final_chain

if __name__=="__main__":    
    question = "你好啊"
    final_chain = get_LECL_chain()
    print('正在回答')
    res = final_chain.invoke(question)
    print(res)

    
    